from collections.abc import Iterable   # Python 3.3+
import numpy as np
import shutil, os
def is_iterable(obj):
    return isinstance(obj, Iterable)

def get_dot_graph(var, verbose=False):
    def _dot_var(var, verbose=False):
        dot_txt = '{} [label="{}", color=orange, style=filled]\n'
        name = ' ' if var.name is None else var.name
        if verbose:
            name += ': '
            name += str(var.shape) + ' ' + str(var.dtype)
        
        return dot_txt.format(id(var), name)
    
    def _doc_func(func):
        dot_txt = '{} [label = "{}", color=lightblue, style=filled, shape=box]\n'
        dot_txt = dot_txt.format(id(func), func.__class__.__name__)
        
        dot_link = '{} -> {}\n'
        for var_input in func.inputs:
            dot_txt += dot_link.format(id(var_input), id(func))
        for var_output in func.outputs:
            dot_txt += dot_link.format(id(func), id(var_output()))
        return dot_txt
    
    #生成绘制计算图的dot语言代码
    dot_txt = ''
    funcs = []
    seen_set = set()
    
    funcs.append(var.creator)
    seen_set.add(var.creator)
    
    while funcs:
        func = funcs.pop()
        for var_input in func.inputs:
            dot_txt += _dot_var(var_input, verbose)
            if var_input.creator is not None and var_input.creator not in seen_set:
                funcs.append(var_input.creator)
                seen_set.add(var_input.creator)
        for var_output in func.outputs:
            dot_txt += _dot_var(var_output(), verbose)
        dot_txt += _doc_func(func)
    
    return 'digraph g {\n' + dot_txt + '}'

def plot_graph(var, image_dir='graph', format='png', verbose=False):
    
    dot_graph = get_dot_graph(var, verbose)
    
    dot_path = shutil.which("dot")        # 先在终端环境里找
    if dot_path:
        os.environ["PATH"] += os.pathsep + os.path.dirname(dot_path)
    else:
        raise RuntimeError("dot not found in PATH")
    
    try:
        from graphviz import Source
        Source(dot_graph).render(
            filename=image_dir,
            format=format,
            cleanup=True
        )
    except ImportError:
        print('graphviz is not installed, cannot plot the graph.')